package com.nx.arch.lock.etcd;

import io.etcd.jetcd.ByteSequence;
import io.etcd.jetcd.Client;
import io.etcd.jetcd.Lease;
import io.etcd.jetcd.Lock;
import io.etcd.jetcd.lease.LeaseKeepAliveResponse;
import io.etcd.jetcd.lock.LockResponse;
import io.grpc.stub.StreamObserver;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.nio.charset.StandardCharsets;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.TimeUnit;

/**
 * @author zs
 * @date 2022/11/6 12:23 PM
 */
public class DemoEtcdLock {
    private final static Logger log = LoggerFactory.getLogger(DemoEtcdLock.class);
    private Client client;
    private Lock lockClient;
    private Lease leaseClient;

    private LockState lockState;

    class LockState{
        private String lockKey;
        private String lockPath;
        private String errorMsg;
        private long leaseTTL;
        private long leaseId;
        private boolean lockSuccess;

        public LockState(String lockKey, long leaseTTL) {
            this.lockKey = lockKey;
            this.leaseTTL = leaseTTL;
        }

        public String getLockKey() {
            return lockKey;
        }

        public void setLockKey(String lockKey) {
            this.lockKey = lockKey;
        }

        public String getLockPath() {
            return lockPath;
        }

        public void setLockPath(String lockPath) {
            this.lockPath = lockPath;
        }

        public String getErrorMsg() {
            return errorMsg;
        }

        public void setErrorMsg(String errorMsg) {
            this.errorMsg = errorMsg;
        }

        public long getLeaseId() {
            return leaseId;
        }

        public void setLeaseId(long leaseId) {
            this.leaseId = leaseId;
        }

        public boolean isLockSuccess() {
            return lockSuccess;
        }

        public void setLockSuccess(boolean lockSuccess) {
            this.lockSuccess = lockSuccess;
        }

        public long getLeaseTTL() {
            return leaseTTL;
        }

        public void setLeaseTTL(long leaseTTL) {
            this.leaseTTL = leaseTTL;
        }
    }


    public DemoEtcdLock(Client client, String lockKey, Long leaseTTL, TimeUnit unit) {
        this.client = client;
        //1.准备客户端
        this.lockClient = client.getLockClient();
        this.leaseClient = client.getLeaseClient();
        this.lockState = new LockState(lockKey,unit.toSeconds(leaseTTL));
    }


    public boolean lock()  {
        try {
            //2.创建租约,并自动续约
            createLease();

            //3.执行加锁，并为锁对应的Key绑定租约
            createLock();
        }catch (InterruptedException | ExecutionException e) {
            //todo:异常处理
        }
        return lockState.isLockSuccess();
    }

    public void unlock()  {
        try {
        //正常释放锁
            if (this.lockState.getLockPath() != null) {
                lockClient.unlock(ByteSequence.from(lockState.getLockPath().getBytes())).get();
            }
            //如果是主动续约，则关闭续约的定时任务

            //删除租约
            if (lockState.getLeaseId() != 0L) {
                leaseClient.revoke(lockState.getLeaseId());
            }
        } catch (InterruptedException | ExecutionException e) {
            //todo:异常处理
        }
        log.info("线程：{} 释放锁", Thread.currentThread().getName());
    }

    // 创建一个租约
    private void createLease() throws ExecutionException, InterruptedException {
        log.debug("[etcd-lock]: start to createLease." + this.lockState.getLockKey() + Thread.currentThread().getName());
        try {
            long leaseId = leaseClient.grant(this.lockState.getLeaseTTL()).get().getID();
            lockState.setLeaseId(leaseId);
            //自动续约
            StreamObserver<LeaseKeepAliveResponse> observer = new StreamObserver<LeaseKeepAliveResponse>() {
                @Override
                public void onNext(LeaseKeepAliveResponse value) {
                    log.trace("cluster node lease remaining ttl: {}, lease id: {}", value.getTTL(), value.getID());
                }

                @Override
                public void onError(Throwable t) {
                    log.error("cluster node lease keep alive failed. exception info: {}", t);
                }

                @Override
                public void onCompleted() {
                    log.trace("cluster node lease completed");
                }
            };
            // 设置自动续约
            leaseClient.keepAlive(leaseId, observer);
        }catch (InterruptedException | ExecutionException e) {
            log.error("[etcd-lock] Create lease failed:" + e);
            lockState.setErrorMsg("Create lease failed:" + e);
            throw e;
        }
    }

    private void createLock() throws ExecutionException, InterruptedException {
        String lockKey = this.lockState.getLockKey();
        log.debug("[etcd-lock]: start to createLock." + lockKey + Thread.currentThread().getName());
        try {
            LockResponse lockResponse = lockClient.lock(ByteSequence.from(lockKey.getBytes()), lockState.getLeaseId()).get();
            if (lockResponse != null) {
                String lockPath = lockResponse.getKey().toString(StandardCharsets.UTF_8);
                this.lockState.setLockPath(lockPath);
                log.info("线程：{} 加锁成功，锁路径：{}", Thread.currentThread().getName(), lockPath);
                this.lockState.setLockSuccess(true);
            }
        }
        catch (InterruptedException | ExecutionException e) {
            log.error("[etcd-lock] lock failed:" + e);
            lockState.setErrorMsg("[etcd-lock] lock failed:" + e);
            leaseClient.revoke(this.lockState.getLeaseId());
            throw e;
        }
    }
}

